using System;
using System.Reflection;
using NBody;
using NBody.Analysis;
using OptionParser;
using DataFile;

[assembly: AssemblyTitle ("NBodyEnergy")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Compute some energy quantities.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class EnergyOptions : InputOutputGetOptions
{
    [Option ("-b", "Number of bins")]
    public int NumBins;
     
    [Option ("-c", "Adjust for centre of mass")]
    public bool CentreOfMass;
    
    [Option ("-e", "Softening length")]
    public double Softening;
    
    [Option ("-f", "File for potential data")]
    public string PotFile;
    
    [Option ("-m", "Output dM/dE")]
    public bool dMdE;
    
    [Option ("-g", "Output the density of states, g(E)")]
    public bool DofS;
    
    [Option ("-d", "Output the distribution function, f(E)")]
    public bool DF;
    
    [Option ("-I", "Assume isotropic orbits, rather than radial")]
    public bool Isotropic;
    
    [Option ("-l", "Output is log(-x), log(y)")]
    public bool Logplot;
    
    [Option ("-s", "Output is x, log(y)")]
    public bool SemiLogplot;
    
    [Option ("-r", "Output total energy versus radius")]
    public bool EnergyVRadius;
    
    public EnergyOptions()
    {
        NumBins = 100;
        Softening = 0.0;
        PotFile = "";
    }
}

public class NBodyEnergy
{
    public static void Main(string[] args)
    {
        EnergyOptions opts = new EnergyOptions();
        opts.ProcessArgs(args, ".energy");
                        
        NBodySystem s = NBodySystem.Read(opts.InFile);
        
        if (opts.CentreOfMass)
            s.AdjustForCentreOfMass(1e-3);
        
        if (opts.PotFile != "")
            Energy.SetPotentialFromData(ref s, Table.Read(opts.PotFile));
        else
            Energy.CalculatePotentialShell(ref s, opts.Softening);
        
        
        if (opts.EnergyVRadius)
        {
            Table data = new Table(s.NumParts, 2, "radius, energy");
            int c = 0;
            double r200 = NBody.Analysis.VirialRadius.CalculateR200(s, NBody.Cosmology.Model.SCDM);
            double beta = 3.69;
            double E0 = -4.0 * Math.PI * beta;
            foreach (EnergyParticle p in s)
            {
                double v = p.AbsoluteVelocity;
                double ke = 0.5 * v * v;
                data[c, 0] = Math.Log10(p.Radius / r200);
                double E = ke + p.Potential;
                data[c, 1] = (-E + E0) / E0 + 1.0;
                c++;
            }
            data.TrimTo(10000);
            data.Print(opts.OutFile);
            Environment.Exit(0);
        }
        
        EnergyBin energy;
        if (opts.dMdE)
            energy = new EnergyBinDMDE();
        else if (opts.DofS)
            energy = new EnergyBinDOS();
        else //if (opts.DF)
            energy = new EnergyBinDF();
            
        if (opts.Isotropic) energy.BinOrbit = EnergyBin.Orbit.Isotropic;
        
        energy.NumBins = opts.NumBins;
        energy.Analyse(s);
        
        if (opts.Logplot)
            energy.Data.Print(opts.OutFile, new Transform[] { MinusLog, Math.Log10 } );
        else if (opts.SemiLogplot)
            energy.Data.Print(opts.OutFile, new Transform[] { Table.IdentityTransform, Math.Log10 } );
        else
            energy.Data.Print(opts.OutFile); 
    }
    
    public static double MinusLog(double x)
    {
        return Math.Log10(-x);
    }
}
